import { Color, Texture2D, Vec3, Node, warn } from "cc";
import { IDisposable } from "../Utils/IDisposable";
import { BuiltInSensorType, IBuiltInSensor } from "./IBuiltInSensor";
import { ISensor } from "./ISensor";
import { CompressionSpec, SensorCompressionType } from "./CompressionSpec";
import { ObservationSpec } from "./ObservationSpec";
import { IGridPerception } from "./IGridPerception";
import { CCCAgentsError } from "../CCCAgentsError";
import { Texture2DWrapper } from "../Utils/Texture2DWrapper";
import { ObservationWriter } from "./ObservationWriter";
import { Utilities } from "../Utilities";


/**
 * @en
 * The way the GridSensor process detected colliders in a cell.
 * @zh
 * GridSensor 处理在单元格中检测到的碰撞器的方式。
 */
export enum ProcessCollidersMethod {

    /**
     * @en
     * Get data from the collider closest to the agent
     * @zh
     *  从距离代理最近的碰撞器获取数据
     */
    ProcessClosestColliders,

    /**
     * @en
     *  Get data from all colliders detected in a cell
     * @zh
     * 从检测到的所有碰撞器中获取数据
     */
    ProcessAllColliders

}

/**
 * @en
 * Grid-based sensor.
 * 
 * @zh
 * 基于网格的传感器。
 */
export class GridSensorBase implements ISensor, IBuiltInSensor, IDisposable {
    private _name: string = "";
    private _cellScale: Vec3 = new Vec3(1, 1, 1);
    private _gridSize: Vec3 = new Vec3(1, 1, 1);
    private _detectableTags: string[] = [];
    public get detectableTags(): string[] {
        return this._detectableTags;
    }
    public set detectableTags(value: string[]) {
        this._detectableTags = value;
    }

    private _compressionType: SensorCompressionType = SensorCompressionType.PNG;
    public get compressionType(): SensorCompressionType {
        return this._compressionType;
    }
    public set compressionType(value: SensorCompressionType) {

        if (!this.isDataNormalized() && value == SensorCompressionType.PNG) {
            warn("GridSensorBase: Data is not normalized, PNG compression will not work correctly.");
            return;
        }

        this._compressionType = value;
    }


    private _observationSpec: ObservationSpec = null!;
     _gridPerception: IGridPerception = null!;
    private _perceptionBuffer: number[] = null!;
    public get perceptionBuffer(): number[] {
        return this._perceptionBuffer;
    }
    public set perceptionBuffer(value: number[]) {
        this._perceptionBuffer = value;
    }

    private _numCells: number = 0;
    private _cellDataBuffer: number[] = null!;
    private _perceptionColors: Color[] = null!;
    private _perceptionTexture: Texture2D = null!;

    private _cellObservationSize: number = 0;


    /**
     * @en
     * Create a GridSensorBase with the specified configuration.
     * 
     * @zh
     * 使用指定的配置创建 GridSensorBase。
     * 
     * @param name the sensor name
     * @param cellScale the scale of each cell in the grid
     * @param gridSize number of cells on each side of the grid
     * @param detectableTags tags to be detected by the sensor
     * @param compressionType compression type
     */
    constructor(name: string,
        cellScale: Vec3,
        gridSize: Vec3,
        detectableTags: string[],
        compressionType: SensorCompressionType,
    ) {
        this._name = name;
        this._cellScale = cellScale;
        this._gridSize = gridSize;
        this._detectableTags = detectableTags;
        this._compressionType = compressionType;

        if (this._gridSize.y != 1) {
            throw new CCCAgentsError("GridSensorBase only supports 2D grids.");
        }

        this._numCells = this._gridSize.x * this._gridSize.z;
        this._cellObservationSize = this.getCellObservationSize();
        this._observationSpec = ObservationSpec.visual(this._gridSize.x, this._gridSize.y, this._detectableTags.length);
        this._perceptionTexture = new Texture2D();
        this._perceptionTexture.reset({
            width: this._gridSize.x,
            height: this._gridSize.z,
            format: Texture2D.PixelFormat.RGBA8888,
            mipmapLevel: 1,
            baseLevel: 0,
            maxLevel: 100
        });

        this.resetPerceptionBuffer();
    }


    public reset(): void { }

    public resetPerceptionBuffer(): void {
        if (this._perceptionBuffer != null) {
            this._cellDataBuffer.fill(0);
            this._perceptionBuffer.fill(0);
        } else {
            this._perceptionBuffer = new Array(this._numCells * this._cellObservationSize);
            this._cellDataBuffer = new Array(this._cellObservationSize);
            this._perceptionColors = new Array(this._numCells);
        }
    }

    public getName(): string {
        return this._name;
    }

    public getCompressionSpec(): CompressionSpec {
        return new CompressionSpec(this._compressionType);
    }

    public getBuiltInSensorType(): BuiltInSensorType {
        return BuiltInSensorType.GridSensor;
    }

    public getCompressedObservation(): Uint8Array {
        let allBytes: Uint8Array = new Uint8Array();
        const numImages = Math.ceil((this._cellObservationSize + 2) / 3);
        for (let i = 0; i < numImages; i++) {
            const channelIndex = i * 3;
            this.gridValuesToTexture(channelIndex, Math.min(3, this._cellObservationSize - channelIndex));
            const bytes = Texture2DWrapper.encodeToPNG(this._perceptionTexture);
            if (bytes) {
                const newAllBytes = new Uint8Array(allBytes.length + bytes.length);
                newAllBytes.set(allBytes);
                newAllBytes.set(bytes, allBytes.length);
                allBytes = newAllBytes;
            }
        }
        return allBytes;
    }

    /**
     * @en
     * Convert observation values to texture for PNG compression.
     * 
     * @zh
     * 将观察值转换为纹理以进行 PNG 压缩。
     * @param channelIndex 
     * @param numChannelsToAdd 
     */
    private gridValuesToTexture(channelIndex: number, numChannelsToAdd: number): void {
        for (let i = 0; i < this._numCells; i++) {
            for (let j = 0; j < numChannelsToAdd; j++) {
                const value = this._perceptionBuffer[i * this._cellObservationSize + channelIndex + j];
                const color = this._perceptionColors[i];
                this.setColorProperties(color, channelIndex + j, value);
            }
        }

        const data = new Uint8Array(this._perceptionColors.length * 4);
        for (let i = 0; i < this._perceptionColors.length; i++) {
            const color = this._perceptionColors[i];
            data[i * 4] = color.r;
            data[i * 4 + 1] = color.g;
            data[i * 4 + 2] = color.b;
            data[i * 4 + 3] = color.a;
        }

        this._perceptionTexture.uploadData(data);
    }

    private setColorProperties(color: Color, index: number, value: number): void {
        if (index === 0) {
            color.r = value;
        } else if (index === 1) {
            color.g = value;
        } else if (index === 2) {
            color.b = value;
        } else if (index === 3) {
            color.a = value;
        }
    }


    /**
     * 
     * @en
     * Get the observation values of the detected game object.
     * Default is to record the detected tag index.
     *
     * This method can be overridden to encode the observation differently or get custom data from the object.
     * When overriding this method, {@link getCellObservationSize} and {@link isDataNormalized}
     * might also need to change accordingly.
     * 
     * @zh
     * 获取检测到的游戏对象的观察值。
     * 默认是记录检测到的标签索引。
     * 
     * 可以重写此方法以不同方式编码观察值或从对象中获取自定义数据。
     * 当重写此方法时，可能还需要相应地更改 {@link getCellObservationSize} 和 {@link isDataNormalized}。
     * 
     * 
     * @example
     * ```ts
     * 
     * protected getObjectData(detectedObject: Node, tagIndex: number, dataBuffer: number[]): void {
     * 
     *      const rigidBody = detectedObject.getComponent(RigidBody);
     *      if (rigidBody) {
     *          const out = new Vec3();
     *          rigidBody.getLinearVelocity(out);
     *          dataBuffer[0] = out.x;
     *          dataBuffer[1] = out.y
     *          dataBuffer[2] = out.z
     *      } else {
     *          dataBuffer[0] = tagIndex + 1;
     *      }
     * }
     * 
     * ```
     * 
     * @param detectedObject The game object that was detected within a certain cell 
     * @param tagIndex The index of the detectedObject's tag in the DetectableObjects list
     * @param dataBuffer The buffer to write the observation values.
     */
    protected getObjectData(detectedObject: Node, tagIndex: number, dataBuffer: number[]): void {
        dataBuffer[0] = tagIndex + 1;
    }

    /**
     * @en
     * Get the observation size for each cell. This will be the size of dataBuffer for {@link "getObjectData}.
     * If overriding {@link "getObjectData}, override this method as well to the custom observation size.
     *
     * @zh
     * 获取每个单元格的观察大小。这将是 {@link "getObjectData"} 的 dataBuffer 的大小。
     * 如果重写 {@link "getObjectData"}，也要相应地重写此方法以获取自定义观察大小。
     * @returns 
     */
    protected getCellObservationSize(): number {
        return 1;
    }

    /**
     * @en
     * Whether the data is normalized within [0, 1]. The sensor can only use PNG compression if the data is normailzed.
     * If overriding {@link getObjectData}, override this method as well according to the custom observation values.
     * 
     * @zh
     * 数据是否在 [0, 1] 内归一化。传感器只能在数据归一化时使用 PNG 压缩。
     * 如果重写 {@link GetObjectData}，也要根据自定义观察值相应地重写此方法。
     * @returns 
     */
    protected isDataNormalized(): boolean {
        return false;
    }

    /**
     * @en
     * Whether to process all detected colliders in a cell. Default to false and only use the one closest to the agent.
     * If overriding {@link getObjectData}, consider override this method when needed.
     * @zh
     * 是否处理单元格中检测到的所有碰撞器。默认为 false，只使用距离代理最近的碰撞器。
     * 
     * @returns 
     */
    getProcessCollidersMethod(): ProcessCollidersMethod {
        return ProcessCollidersMethod.ProcessClosestColliders;
    }

    /**
     * @en
     * If using PNG compression, check if the values are normalized.
     * 
     * @zh
     * 如果使用 PNG 压缩，请检查值是否已归一化。
     * 
     * @param dataValues 
     * @param detectedObject 
     * @returns 
     */
    private validateValues(dataValues: number[], detectedObject: Node): void {
        if (this._compressionType !== SensorCompressionType.PNG) {
            return;
        }

        for (let j = 0; j < dataValues.length; j++) {
            if (dataValues[j] < 0 || dataValues[j] > 1) {
                throw new CCCAgentsError(`When using compression type ${this._compressionType} the data value has to be normalized between 0-1.
                 Received value[${dataValues[j]}] for ${detectedObject.name}`);
            }
        }
    }

    /**
     * @en
     *  Collect data from the detected object if a detectable tag is matched.
     * 
     * @zh
     * 如果匹配了可检测标签，则从检测到的对象收集数据。
     * 
     * @param detectedObject 
     * @param cellIndex 
     */
    processDetectedObject(detectedObject: Node, cellIndex: number): void {
        for (let i = 0; i < this._detectableTags.length; i++) {
            if (detectedObject != null && detectedObject.name === (this._detectableTags[i])) {
                if (this.getProcessCollidersMethod() === ProcessCollidersMethod.ProcessAllColliders) {
                    this._cellDataBuffer.copyWithin(0, cellIndex * this._cellObservationSize, cellIndex * this._cellObservationSize + this._cellObservationSize);

                    this._cellDataBuffer = this._perceptionBuffer.slice(cellIndex * this._cellObservationSize, cellIndex * this._cellObservationSize + this._cellObservationSize);
                } else {
                    this._cellDataBuffer.fill(0);
                }

                this.getObjectData(detectedObject, i, this._cellDataBuffer);
                this.validateValues(this._cellDataBuffer, detectedObject);
                this._perceptionBuffer.splice(cellIndex * this._cellObservationSize, this._cellObservationSize, ...this._cellDataBuffer);
                break;
            }
        }
    }


    public update(): void {
        this.resetPerceptionBuffer();
        if (this._gridPerception != null) {
            this._gridPerception.perceive();
        }
    }

    public getObservationSpec(): ObservationSpec {
        return this._observationSpec;
    }

    public write(writer: ObservationWriter): number {
        let index = 0;
        for (let h = this._gridSize.z - 1; h >= 0; h--) {
            for (let w = 0; w < this._gridSize.x; w++) {
                for (let d = 0; d < this._cellObservationSize; d++) {
                    writer.setItem3D(h, w, d, this._perceptionBuffer[index]);
                    index++;
                }
            }
        }
        return index;
    }

    public dispose(): void {
        if (this._perceptionTexture != null) {
            Utilities.destroyTexture(this._perceptionTexture);
            this._perceptionTexture = null!;
        }
    }
}